import os
from PIL import Image
import numpy as np


def parse_annotation_file(annotation_file, image_width, image_height):
    with open(annotation_file, 'r') as f:
        lines = f.readlines()
        annotations = []
        for line in lines:
            values = line.split()
            category_id = int(values[0])  # Extract category id
            # Convert normalized coordinates to absolute coordinates
            # x_center = float(values[1]) * image_width
            # y_center = float(values[2]) * image_height
            # width = float(values[3]) * image_width
            # height = float(values[4]) * image_height

            # Calculate the coordinates of the four vertices
            # x1 = x_center - width / 2
            # y1 = y_center - height / 2
            # x2 = x_center + width / 2
            # y2 = y_center - height / 2
            # x3 = x_center + width / 2
            # y3 = y_center + height / 2
            # x4 = x_center - width / 2
            # y4 = y_center + height / 2
            x1 = float(values[1])
            y1 = float(values[2])
            x2 = float(values[3])
            y2 = float(values[4])
            x3 = float(values[5])
            y3 = float(values[6])
            x4 = float(values[7])
            y4 = float(values[8])
            normalized_coords = [x1,y1,x2,y2,x3,y3,x4,y4]
            normalized_coords = np.array(normalized_coords).reshape(-1, 2)
            print(normalized_coords)

            # Convert to pixel coordinates
            pixel_coords = normalized_coords * [image_width, image_height]

            pixel_coords = np.maximum(pixel_coords, 0)

            # Append the coordinates to the annotations list
            list=pixel_coords.flatten().tolist()
            print(type(list))
            print(list)
            annotations.append((category_id, list[0],list[1],list[2],list[3],list[4],list[5],list[6],list[7]))
        return annotations


def process_annotation_files(image_folder, annotation_folder):
    for annotation_file in os.listdir(annotation_folder):
        if annotation_file.endswith('.txt'):
            # Check for different image formats
            image_filename = os.path.splitext(annotation_file)[0]
            image_path = None
            for ext in ['.jpeg', '.jpg', '.JPEG', '.JPG', '.png', '.PNG']:
                potential_path = os.path.join(image_folder, image_filename + ext)
                if os.path.exists(potential_path):
                    image_path = potential_path
                    break

            if image_path is None:
                print(f"No corresponding image found for {annotation_file}")
                continue

            # Get image dimensions
            image = Image.open(image_path)
            image_width, image_height = image.size
            print(image_width, image_height)

            # Parse the annotation file
            annotation_path = os.path.join(annotation_folder, annotation_file)
            annotations = parse_annotation_file(annotation_path, image_width, image_height)

            # Write the transformed annotations to a new file
            output_filename = os.path.splitext(annotation_file)[0] + '_transformed.txt'
            output_path = os.path.join(annotation_folder, output_filename)
            with open(output_path, 'w') as f:
                for annotation in annotations:
                    f.write(' '.join(str(coord) for coord in annotation) + '\n')


# def process_annotation_files(image_folder, annotation_folder):
#     for annotation_file in os.listdir(annotation_folder):
#         if annotation_file.endswith('.txt'):
#             # Get the corresponding image filename
#             image_filename = os.path.splitext(annotation_file)[0] + '.jpeg'
#             image_path = os.path.join(image_folder, image_filename)
#
#             # Get image dimensions
#             image = Image.open(image_path)
#             image_width, image_height = image.size
#             print(image_width,image_height)
#
#             # Parse the annotation file
#             annotation_path = os.path.join(annotation_folder, annotation_file)
#             annotations = parse_annotation_file(annotation_path, image_width, image_height)
#
#             # Write the transformed annotations to a new file
#             output_filename = os.path.splitext(annotation_file)[0] + '_transformed.txt'
#             output_path = os.path.join(annotation_folder, output_filename)
#             with open(output_path, 'w') as f:
#                 for annotation in annotations:
#                     f.write(' '.join(str(coord) for coord in annotation) + '\n')
# # Example usage:
# image_folder = r'D:\Work\PythonProjects\wyr\ultralytics-main\runs\obb\predict23'
# annotation_folder = r'D:\Work\PythonProjects\wyr\ultralytics-main\runs\obb\predict23\labels'
# process_annotation_files(image_folder, annotation_folder)